#ifndef _I_SPATIAL_TREE_H_
#define _I_SPATIAL_TREE_H_

#include "../cuCommon/cuTypes.h"

#include "../interface/I_Buffer.h"
#include "../interface/I_DoubleBuffer.h"

#include "../rtCommon/Box.h"

#define X											0x00
#define Y											0x01
#define Z											0x02
#define XYZ										0x03

#define TREE_AXIS_X					   0x00
#define TREE_AXIS_Y					   0x01
#define TREE_AXIS_Z					   0x02
#define TREE_AXIS_W					   0x03
#define TREE_AXIS_XYZ				   0x03

#define TREE_LEAF              0x03
#define TREE_NON_LEAF          0x00

#define TREE_ROOT              0x00

#define KD_DEFAULT_BIN_EXP   0x0A

//////////////////////////////////////////////////////////////////////////
/// S_Interval structure 
//////////////////////////////////////////////////////////////////////////
struct S_Interval
{
	float start, end;

	S_Interval() { start = 0; end = 0; }
	S_Interval(float l, float r) { start = l; end = r; }

	S_Interval & operator = (const float2     &i) { start = i.x;     end = i.y;   return *this; }
	S_Interval & operator = (const S_Interval &i) { start = i.start; end = i.end; return *this; }

	operator float2() { return f2(start, end); }

	float length() { return end - start; }
};

//////////////////////////////////////////////////////////////////////////
/// Tree Node structure used during build stage (aligned to 128 bits)
//////////////////////////////////////////////////////////////////////////
struct S_TreeNode
{
	//////////////////////////////////////////////////////////////////////////
	/// split position with axis coded in lower bits
	//////////////////////////////////////////////////////////////////////////
	union
	{
		float split; // store splitting position
		int   bin;   // or bin index
	};

	//////////////////////////////////////////////////////////////////////////
	/// Start Index / Left Node union
	//////////////////////////////////////////////////////////////////////////
	union 
	{
		int  start;  
		int  left;
	};

	//////////////////////////////////////////////////////////////////////////
	/// Start Index 
	//////////////////////////////////////////////////////////////////////////
	union
	{
		int  end;  
		int  right;
	};

	//////////////////////////////////////////////////////////////////////////
	/// parent index to build traversal tree
	//////////////////////////////////////////////////////////////////////////
	union 
	{
		uint32 parent;
		float  sah;
	};

	S_TreeNode()
	{
		split = as_float(TREE_NON_LEAF);
		start = 0;
		end   = 0;
		parent = TREE_ROOT;
	}

	S_TreeNode(float s)
	{
		split = s;
		start = 0;
		end   = 0;
		parent = TREE_ROOT;
	}

	S_TreeNode(int p)
	{
		split = as_float(TREE_NON_LEAF);
		start = 0;
		end   = 0;
		parent = p;
	}

	S_TreeNode(float s, int p)
	{
		split = s;
		start = 0;
		end   = 0;
		parent = p;
	}

	S_TreeNode(int s, int e)
	{
		split = as_float(TREE_NON_LEAF);
		start = s;
		end   = e;
		parent = TREE_ROOT;
	}

	S_TreeNode(int nStart, int nEnd, int nParent)
	{
		start = nStart;
		end = nEnd;
		parent = nParent;
		if (end - start == 0) SplitAxis(0.0f, TREE_LEAF);
	}

	inline bool		Leaf()	{ return (as_uint(split) & TREE_LEAF) == TREE_LEAF;	}
	inline void		SplitAxis(float s, int axis)	{ split = as_float(as_uint(s) & (~TREE_LEAF) | axis); }
	inline float	Split()	{ return as_float(as_uint(split) & (~TREE_LEAF)); }
	inline int		Axis()	{ return (as_uint(split) &  TREE_LEAF ); }
	inline int		Count()	{ return end - start; }
};

//////////////////////////////////////////////////////////////////////////
/// KdNode structure used for KdTree
//////////////////////////////////////////////////////////////////////////
struct S_KdNode
{
	union 
	{
		float    split;      // for non leaf nodes TN holds split position with axis encoded in lower bits
		int      count;      // for leaf node there is no split, but there is primitive count stores shifted right by two
		float    t;          // for tracing t stores parametric intersection data.
	};

	union 
	{
		int   right;      // for non leaf nodes TN holds offset to the right sub-node
		int   start;      // for leaf node there is no right sub-node, but there is an index into primIdx soup
		int   next;       // for tracing next stores stacked index
	};

	S_KdNode()
	{
		split = 0;
		right = 0;
	}
	
	S_KdNode(int node_index, float tfar)
	{
		next = node_index;
		t = tfar;
	}

	inline void		SplitAxis(float s, int axis)	{ split = as_float(as_uint(s) & (~TREE_LEAF) | axis); }
	inline float	Split()	{ return as_float(as_uint(split) & (~TREE_LEAF)); }
	inline int		Axis()	{ return (count & TREE_LEAF); }
	inline int		Count()	{ return (count >> 2); }
	inline int		Start() { return start; }
};


//////////////////////////////////////////////////////////////////////////
/// BVHNode structure used for BVHTree
//////////////////////////////////////////////////////////////////////////
struct S_BVHNode
{
	union 
	{
		struct
		{
			float3 pmin;
			union
			{
				int right;
				int start;
			};
		};
		float4 pmin_rs;
	};

	union
	{
		struct
		{
			float3 pmax;
			union
			{
				float  split;
				int axis;
				int count;
			};
		};
		float4 pmax_lc;
	};

	inline int		Start() { return start; }
	inline void		SplitAxis(float s, int axis)	{ split = as_float(as_uint(s) & (~TREE_LEAF) | axis); }
	inline int		Axis()	{ return (axis & TREE_LEAF ); }
	inline int		Count()	{ return (count >> 2); }
	inline float	Split()	{ return as_float(as_uint(split) & (~TREE_LEAF)); }
};

struct S_BVHNode_cu
{
	float4 pmin_rs;
	float4 pmax_lc;
};

enum e_SpatialTree_Type
{
	e_SpatialTree_Kd = 64,
	e_SpatialTree_Bvh = 128,
};

enum e_SpatialTree_Data
{
	e_SpatialTree_Data_PrimitiveBI = 1,
	e_SpatialTree_Data_WeightStart = 2,
	e_SpatialTree_Data_PrimIdx     = 4,
	e_SpatialTree_Data_TreeIdx     = 8,
	e_SpatialTree_Data_Node        = 16,
	e_SpatialTree_Data_Interval    = 32,
	e_SpatialTree_Data_Kd          = 64,
	e_SpatialTree_Data_BVH         = 128,
};

//////////////////////////////////////////////////////////////////////////
/// Spatial data structure interface
//////////////////////////////////////////////////////////////////////////
class I_SpatialTree
{
	friend class I_Asset;
	friend class I_SceneGraph;
	friend class C_gl_SpatialTree;

protected:
	//////////////////////////////////////////////////////////////////////////
	/// Spatial tree doesnt have its data but it stores data pointers from the asset
	//////////////////////////////////////////////////////////////////////////
	I_Buffer					*	m_pVertex;											// pointer to position packed in f4
	int									m_nVertex;
	I_Buffer					*	m_pIndex;												// pointer to regular position index
	int									m_nPrimitive;
	
	//////////////////////////////////////////////////////////////////////////
	/// If Spatial tree holds actors ActorData holds per actor info
	/// x. 
	/// y.
	/// z.
	/// w.
	//////////////////////////////////////////////////////////////////////////
	I_Buffer					*	m_pActorData;

	//////////////////////////////////////////////////////////////////////////
	/// shared data - may actors store data in these arrays with offsets
	//////////////////////////////////////////////////////////////////////////
	I_Buffer					*	m_pTreeIdx;											// large index buffer that is used for tree construction - stores results
	int									m_nTreeIdx;											// m_nTree is used as an index into m_pTreeIdx
	I_Buffer					*	m_pKdTree;											// KdTree is a packed array of f2 that can be cast to S_KdNode
	I_Buffer					*	m_pBVHTree;											// BVH    is a packed array of f8 that can be cast to S_BVHNode

	//////////////////////////////////////////////////////////////////////////
	/// common data - many actors can reuse these arrays
	//////////////////////////////////////////////////////////////////////////
	I_DoubleBuffer		*	m_pPrimIdx;											// large index buffer that is used for tree construction - stores intermediate data
	int									m_nPrimIdx;											// m_nPrimitive is used as an index into m_primIdx arrays

	I_Buffer					*	m_ppBoundingInterval[XYZ];						// per primitive bounding interval values packed in 3 different axis arrays

	I_Buffer					*	m_pNode;												// m_Tree is an internal data structure that couples with m_Interval to construct 
	int									m_nNode;

	I_Buffer					*	m_pInterval;										// various data structures such as KdTree or BVH
	int									m_nInterval;
 
	//////////////////////////////////////////////////////////////////////////
	/// 
	//////////////////////////////////////////////////////////////////////////
	S_Box								m_Box;													// Bounding box for the whole scene

	int									m_nDepth;
	int									m_nMaxPrim;											// Maximum primitives inside a node

	int									m_nTraversalIdx;

	void								Null();

public:

	I_SpatialTree();
	~I_SpatialTree();

	//////////////////////////////////////////////////////////////////////////
	/// Create function allocates necessary resources
	//////////////////////////////////////////////////////////////////////////
	virtual	void								Create(int nFlag, 
																		I_Buffer *pVertex, 
																		I_Buffer *pIndex,
																		int nVertex, int nPrimitive, 
																		int nCount, int nDepth,
																		I_api_Device * pDevice=NULL)=0;

	virtual	void								Create(int nFlag, 
																		I_Buffer *pVertex, 
																		I_Buffer *pIndex,
																		int nVertex, int nPrimitive, 
																		int nCount, int nDepth,
																		I_SpatialTree * pShared)=0;

	//////////////////////////////////////////////////////////////////////////
	/// Release resources
	//////////////////////////////////////////////////////////////////////////
	void								Release();

	//////////////////////////////////////////////////////////////////////////
	/// 1. Compute per primitive AABB and store them in split[XYZ]
	//////////////////////////////////////////////////////////////////////////
	virtual void				PrimitiveAABB()=0;

	//////////////////////////////////////////////////////////////////////////
	/// KdTree_p0 function creates KdTree Tree data structure using SAH
	//////////////////////////////////////////////////////////////////////////
	virtual void				KdTree_p0()=0;
	virtual	void				KdTree_Instanced_p0()=0;

	//////////////////////////////////////////////////////////////////////////
	/// KdTree_p1 and p2 creates a KdTree data structure from Tree recursively
	//////////////////////////////////////////////////////////////////////////
	void								KdTree_p1();
	void								KdTree_p2(S_TreeNode *pb, S_KdNode * pt, S_TreeNode & n);

	//////////////////////////////////////////////////////////////////////////
	/// BVHTree_p0 function creates Tree data structure using SAH
	//////////////////////////////////////////////////////////////////////////
	virtual void				BVHTree_p0()=0;
	virtual	void				BVHTree_Instanced_p0()=0;

	//////////////////////////////////////////////////////////////////////////
	/// BVHTree_p1 creates a BVH data structure from Tree recursively
	//////////////////////////////////////////////////////////////////////////
	void								BVHTree_p1();
	void								BVHTree_p2(S_TreeNode *pb, S_BVHNode * pt, S_Interval * pi, S_TreeNode & n, int i);

	//////////////////////////////////////////////////////////////////////////
	/// Get TraversalData
	//////////////////////////////////////////////////////////////////////////
	I_Buffer	*					KdTree()		{ return m_pKdTree;	  }
	I_Buffer	*					BVHTree()		{ return m_pBVHTree;  }
	I_Buffer	*					Interval()	{ return m_pInterval; }
	I_Buffer	*					Node()			{ return m_pNode;		  }
	I_Buffer	*					TreeIdx()		{ return m_pTreeIdx;  }
	I_Buffer	*					Vertex()		{ return m_pVertex;	  }
	I_Buffer	*					Primitive()	{ return m_pIndex;	  }
	I_Buffer	*					ActorData() { return m_pActorData; }
	I_DoubleBuffer	*		PrimIdx()		{ return m_pPrimIdx;  }
	I_Buffer	*					PrimIdx(int buffer)	{ return m_pPrimIdx->Buffer(buffer);  }

	I_Buffer	*					PrimitiveBI(int nAxis) { return m_ppBoundingInterval[nAxis]; }

	virtual	void				KdTree(I_api_Device * pDevice, int nMaxDepth=20, int nCount=1)=0;
	virtual	void				BVHTree(I_api_Device * pDevice, int nMaxDepth=20, int nCount=1)=0;

	virtual	void				KdTree(I_SpatialTree * pShared, int nMaxDepth=20, int nCount=1)=0;
	virtual	void				BVHTree(I_SpatialTree * pShared, int nMaxDepth=20, int nCount=1)=0;

	virtual	void				ActorData(int nCount, I_api_Device *pDevice)=0;
	
	int									nPrimitive()	{ return m_nPrimitive; }
	int									nVertex()			{ return m_nVertex; }
	int									nDepth()			{ return m_nDepth; }

	//////////////////////////////////////////////////////////////////////////
	/// Get MaxDepth
	//////////////////////////////////////////////////////////////////////////
	int32								maxNode() { return m_nTraversalIdx; }

	//////////////////////////////////////////////////////////////////////////
	/// Get SAABox
	//////////////////////////////////////////////////////////////////////////
	S_Box								Box() { return m_Box; }
	void								Box(const S_Box & b) { m_Box = b; }

	//////////////////////////////////////////////////////////////////////////
	/// Render methods
	//////////////////////////////////////////////////////////////////////////
	//virtual void RenderKdTree(int maxLevel);
	virtual void				RenderBVHTree(int maxLevel, matrix4x4 * pModel, matrix4x4 * pView, matrix4x4 * pProjection, I_api_Device * pApiDevice)=0;
	virtual void				RenderBVHTreeRayTrace(float3 orig, float3 dir, int maxLevel, bool empty)=0;
};

#endif _I_SPATIAL_TREE_H_